import os

from matplotlib import pyplot as plt
from matplotlib.ticker import MaxNLocator


def draw_single_picture(x_array, y_array, title=None, x_lable=None, y_lable=None, zoom_size=None, save_path=None,
                        hline=None):
    plt.plot(x_array, y_array)
    if hline:
        plt.hlines(hline, x_array[0], x_array[-1], colors='red', linestyles='dashed')
    if zoom_size:
        plt.axis(zoom_size)
    plt.title(title)
    plt.xlabel(x_lable)
    plt.ylabel(y_lable)
    plt.yticks([3.25,5.81])
    if save_path:
        plt.savefig(save_path)
    plt.show()


def draw_multiple_picture(x_array_list, y_array_list, title, y_label_array_list, zoom_config=None, save_path=None,
                          hline_config=None):
    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(12, 9))
    # [[ax1,ax2][ax3,ax4]] = [ax1,ax2,ax3,ax4]
    fig.suptitle(title)
    axes_list = [i for j in axes for i in j]
    for ax, x, y, y_label in zip(axes_list, x_array_list, y_array_list, y_label_array_list):
        ax.plot(x, y)
        ax.set_ylabel(y_label)
        ax_index = axes_list.index(ax)
        if hline_config and hline_config.get(ax_index):
            ax.hlines(hline_config.get(ax_index), x[0], x[-1], colors='red', linestyles='dotted')
        if zoom_config and zoom_config.get(ax_index):
            ax.set_ylim(*zoom_config.get(ax_index))
    plt.tight_layout()
    if save_path:
        plt.savefig(save_path)
    plt.show()


def draw_multiple_data_picture(x_array_list, y_array_list, title, y_label_array_list, zoom_config=None, save_path=None,
                               hline_config=None, nrows=1, ncols=1, legend_config=None):
    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(12, 9))
    fig.suptitle(title)
    if ncols + nrows > 3:
        axes_list = [i for j in axes for i in j]
    else:
        axes_list = [axes]
    for ax, x, y, y_label in zip(axes_list, x_array_list, y_array_list, y_label_array_list):

        if isinstance(y, list):
            for sub_y in y:
                ax.plot(x, sub_y)
        else:
            ax.plot(x, y)
        y_tick_list = ax.get_yticks()
        ax.set_yticks(y_tick_list[::int(len(y_tick_list)/5)])
        ax.set_ylabel(y_label)
        ax_index = axes_list.index(ax)
        if hline_config and hline_config.get(ax_index):
            ax.hlines(hline_config.get(ax_index), x[0], x[-1], colors='red', linestyles='dotted')
        if zoom_config and zoom_config.get(ax_index):
            ax.set_ylim(*zoom_config.get(ax_index))
        if legend_config:
            ax.legend(legend_config[ax_index])

    plt.tight_layout()
    if save_path:
        plt.savefig(save_path)
    plt.show()


if __name__ == '__main__':
    draw_single_picture([1, 2, 3], [4, 5, 6], 'title', 'x', 'y', hline=5)
    print(os.getcwd())
